#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/6/15 19:17
# @Author  : LiShan
# @Email   : lishan_1997@126.com
# @File    : lane_segmenr.py
# @Note    : this is note
import os
import cv2.cv2 as cv
import numpy as np

save_path = "./lane"
width = 480
height = 360
if os.path.exists(save_path):
    pass
else:
    os.mkdir(save_path)


def get_distance_from_point_to_line(point, line_point1, line_point2):
    # 对于两点坐标为同一点时,返回点与点的距离
    # if line_point1 == line_point2:
    #     point_array = np.array(point )
    #     point1_array = np.array(line_point1)
    #     return np.linalg.norm(point_array -point1_array )
    # 计算直线的三个参数
    A = line_point2[1] - line_point1[1]
    B = line_point1[0] - line_point2[0]
    C = (line_point1[1] - line_point2[1]) * line_point1[0] + \
        (line_point2[0] - line_point1[0]) * line_point1[1]
    # 根据点到直线的距离公式计算距离
    distance = np.abs(A * point[0] + B * point[1] + C) / (np.sqrt(A ** 2 + B ** 2))
    return distance


def point_distance_line(point, line_point1, line_point2):
    # 计算向量
    vec1 = line_point1 - point
    vec2 = line_point2 - point
    distance = np.abs(np.cross(vec1, vec2)) / np.linalg.norm(line_point1 - line_point2)
    return distance


#
# def extend_line(x1, y1, x2, y2, w, h):
#     print(666)

def max_line(L1, L2):
    # L1：直线1，L2：直线2
    x1, y1, x2, y2 = L1
    x3, y3, x4, y4 = L2
    # 取最高的起点
    if y1 < y3:
        x5, y5 = x1, y1
    else:
        x5, y5 = x3, y3
    # 取最低的终点
    if y2 > y4:
        x6, y6 = x2, y2
    else:
        x6, y6 = x4, y4
    # 返回合并后的直线
    return x5, y5, x6, y6


def adjust_location(lines):
    # 调整直线端点位置，从上到下
    for line in lines:
        if line[1] > line[3]:
            line[0], line[1] = line[2], line[3]
    # 调整直线顺序，从左到右，（下端点判断）
    for i in range(len(lines)):
        for j in range(i + 1, len(lines)):
            if lines[i][2] > lines[j][2]:
                lines[i], lines[j] = lines[j], lines[i]
    return lines


def getkpoints(imag, input1, maxCorners=150, qualityLevel=0.01, minDistance=5):
    mask1 = np.zeros_like(input1)
    x = 0
    y = 0
    w1, h1 = input1.shape
    input1 = input1[0:w1, 0:h1]
    try:
        w, h = imag.shape
    except:
        return None
    mask1[y:y + h, x:x + w] = 255
    keypoints = []
    kp = cv.goodFeaturesToTrack(input1, maxCorners, qualityLevel, minDistance)
    if kp is not None and len(kp) > 0:
        for x, y in np.float32(kp).reshape(-1, 2):
            keypoints.append((x, y))
    return keypoints


def filter_line(lines, mink=0.5, thk=0.1, thd=3):
    lane_line = []
    lane_k = []
    for line in lines:
        x1, y1, x2, y2 = line[0]
        # 求直线斜率
        if x2 != x1:
            k = (y2 - y1) / (x2 - x1)
        else:
            k = (y2 - y1) / (x2 - x1 + np.finfo(float).eps)
        # 排除较为水平的直线（车道线一般为竖直线）
        if np.abs(k) <= mink:
            continue
        else:
            # 合并相近直线
            for i in range(len(lane_k)):
                # 搜索斜率相近的直线
                if np.abs(k - lane_k[i]) < thk:
                    L1 = lane_line[i]
                    L2 = line[0]
                    p0 = np.array([x1, y1])
                    p1 = np.array([L1[0], L1[1]])
                    p2 = np.array([L1[2], L1[3]])
                    # 判断两直线间距是否满足阈值
                    d = point_distance_line(p0, p1, p2)
                    if d < thd:
                        # 更新直线
                        x1, y1, x2, y2 = max_line(L1, L2)
                        lane_line[i] = [x1, y1, x2, y2]
                        # 更新直线斜率
                        if x2 != x1:
                            k = (y2 - y1) / (x2 - x1)
                        else:
                            k = (y2 - y1) / (x2 - x1 + np.finfo(float).eps)
                        lane_k[i] = k
                        break
            # 不存在相近直线，加入新直线
            else:
                lane_k.append(k)
                lane_line.append([x1, y1, x2, y2])
    # 调整直线位置顺序(从上到下，从左到右)
    lane_line = adjust_location(lane_line)
    return lane_line


def std_hough_line(origin, edge, threshold=160):
    # image、rho、theta、threshold
    huogh_lane_pre = origin.copy()
    huogh_lane = origin.copy()
    lines = cv.HoughLines(edge, 1, np.pi / 180, threshold)
    print("标准霍夫变换共预检测出%d条车道线" % len(lines))
    lines_list = []
    for line in lines:
        rho, theta = line[0]
        a = np.cos(theta)
        b = np.sin(theta)
        x0 = a * rho
        y0 = b * rho
        x1 = int(x0 + 1000 * (- b))
        y1 = int(y0 + 1000 * a)
        x2 = int(x0 - 1000 * (- b))
        y2 = int(y0 - 1000 * a)
        lines_list.append([[x1, y1, x2, y2]])
    if len(lines_list):
        for i in range(len(lines_list)):
            x1, y1, x2, y2 = lines_list[i][0]
            cv.line(huogh_lane_pre, (x1, y1), (x2, y2), (0, 0, 255), 2)
        cv.imwrite(save_path + "/std_huogh_lane_pre.jpg", huogh_lane_pre)
    lane_line = filter_line(lines_list)
    print("标准霍夫变换筛选后剩下%d条车道线" % len(lane_line))
    if len(lane_line):
        for i in range(len(lane_line)):
            x1, y1, x2, y2 = lane_line[i]
            cv.line(huogh_lane, (x1, y1), (x2, y2), (0, 0, 255), 2)
        cv.imwrite(save_path + "/std_huogh_lane.jpg", huogh_lane)


def prob_hough_line(origin, edge, threshold=160, Length=200, Gap=100):
    # image、rho、theta、threshold、minLineLength、maxLineGap
    huoghP_lane_pre = origin.copy()
    huoghP_lane = origin.copy()
    lines = cv.HoughLinesP(edge, 1, np.pi / 180, threshold, minLineLength=Length, maxLineGap=Gap)
    print("概率霍夫变换共预检测出%d条车道线" % len(lines))
    if len(lines):
        for i in range(len(lines)):
            x1, y1, x2, y2 = lines[i][0]
            cv.line(huoghP_lane_pre, (x1, y1), (x2, y2), (0, 0, 255), 2)
    cv.imwrite(save_path + "/pro_huoghP_lane_pre.jpg", huoghP_lane_pre)
    lane_line = filter_line(lines)
    color = [(255, 0, 255), (255, 255, 0), (0, 255, 255)]
    color_mask = np.zeros(huoghP_lane.shape, np.uint8)
    bw_mask = np.zeros(huoghP_lane.shape, np.uint8)
    seg_lane = np.zeros(huoghP_lane.shape, np.uint8)
    print("概率霍夫变换筛选后剩下%d条车道线" % len(lane_line))
    if len(lane_line):
        for i in range(len(lane_line)):
            # 当前车道线
            x1, y1, x2, y2 = lane_line[i]
            # 绘制车道线
            cv.line(huoghP_lane, (x1, y1), (x2, y2), (0, 0, 255), 2)
            # 上一条车道线
            if i > 0:
                x3, y3, x4, y4 = lane_line[i - 1]
            else:
                x3, y3, x4, y4 = 0, 0, 0, height
            # 组合车道边框定点
            points = [(x3, y3), (x4, y4), (x2, y2), (x1, y1)]
            pts = np.array([points], np.int32)
            pts = pts.reshape((-1, 1, 2))
            # --------------画车道边框---------------------
            color_mask = cv.polylines(color_mask, [pts], True, (0, 255, 255))
            bw_mask = cv.polylines(bw_mask, [pts], True, (0, 255, 255))
            # -------------填充车道---------------------
            color_mask = cv.fillPoly(color_mask, [pts], color[i % 3])
            bw_mask = cv.fillPoly(bw_mask, [pts], (255, 255, 255))
            seg_lane = cv.bitwise_and(color_mask, huoghP_lane)
    cv.imwrite(save_path + "/color_mask.jpg", color_mask)
    cv.imwrite(save_path + "/bw_mask.jpg", bw_mask)
    cv.imwrite(save_path + "/fill_lane.jpg", seg_lane)
    cv.imwrite(save_path + "/pro_huogh_lane.jpg", huoghP_lane)
    return bw_mask, color_mask


def lane_line_detect(img):
    # 尺寸变化
    img = cv.resize(img, (width, height))
    # 原图
    origin = img
    cv.imwrite(save_path + "/origin.jpg", origin)
    # 灰度化
    gray = cv.cvtColor(origin, cv.COLOR_BGR2GRAY)
    cv.imwrite(save_path + "/gray.jpg", gray)
    # 直方图均衡化
    equ = cv.equalizeHist(gray)
    cv.imwrite(save_path + "/equ.jpg", equ)
    # 二值化
    ret, binary = cv.threshold(gray, 150, 255, cv.THRESH_BINARY)
    cv.imwrite(save_path + "/binary.jpg", binary)
    # 边缘检测
    edge = cv.Canny(gray, 50, 150, apertureSize=3)
    cv.imwrite(save_path + "/edge.jpg", edge)
    # 标准霍夫变换寻找车道线
    std_hough_line(origin, edge, 160)
    # 概率霍夫变换寻找车道线
    bw_mask, color_mask = prob_hough_line(origin, edge, 160)
    return bw_mask, color_mask


# 目标预检测
def pre_object_detect(contours, width, height, ra=None, rw=None, rh=None):
    # ra = [0.03, 0.5]
    # rw = [0.07, 0.3]
    # rh = [0.01, 0.1]
    if ra is None:
        ra = [0.03, 0.5]
    if rw is None:
        rw = [0.08, 0.3]
    if rh is None:
        rh = [0.03, 0.1]
    pre_object = []
    for cnt in contours:
        # 最小外接矩形
        (cx, cy), (w, h), rotate = cv.minAreaRect(cnt)
        # 边框矩形
        # x, y, w, h = cv.boundingRect(cnt)
        area = cv.contourArea(cnt)
        img_area = width * height
        area_rate = area / img_area * 100
        w_rate = w / width
        h_rate = h / height
        # print("面积：%.10f" % area_rate)
        # print("宽度：%.5f" % w_rate)
        # print("高度：%.5f" % h_rate)
        # 过滤掉小物体、大物体
        c1 = area_rate < ra[0] or area_rate > ra[1]
        c2 = w_rate < rw[0] or w_rate > rw[1]
        c3 = h_rate < rh[0] or h_rate > rh[1]
        if c1 or c2 or c3:
            continue
        else:
            pre_object.append(cnt)
    return pre_object


def counter_similarity(cnt, dis=35, simi=0.8):
    remove_cnt = []
    new_cnt = []
    for i in range(len(cnt)):
        for j in range(i, len(cnt)):
            if i == j:
                continue
            else:
                # print("轮廓%d：%d, 轮廓%d：%d" % (i, j, len(cnt[i]), len(cnt[j])))
                num = min(len(cnt[i]), len(cnt[j]))
                sim_num = 0
                for k in range(num):
                    d = sum(np.abs((cnt[i][k] - cnt[j][k])[0]))
                    if d < dis:
                        sim_num += 1
                similarity = (sim_num / num)
                # print("相似度：%.2f" % similarity)
                if similarity >= simi:
                    remove_cnt.append(i)
    for i in range(len(cnt)):
        if i not in remove_cnt:
            new_cnt.append(cnt[i])
    return new_cnt


def templte_read(templte):
    if templte == 0 or templte == "straight":
        # 读取模板
        img = cv.imread("./templte/straight.png")
        # img = cv.resize(img, (width, height))
        # 边缘检测
        edge = cv.Canny(img, 50, 150, apertureSize=3)
        cv.imwrite(save_path + "/straight_edge.jpg", edge)
        # 二值图像提取轮廓，轮廓跟踪算法（Suzuki，1985）
        contours, hierarchy = cv.findContours(edge, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
        contour_origin_pre = img.copy()
        contour_origin = img.copy()
        cnt = contours[0:len(contours)]
        cv.drawContours(contour_origin_pre, cnt, -1, (0, 0, 255), 2)
        cv.imwrite(save_path + "/str_contour_pre.jpg", contour_origin_pre)
        print("预检直行车道地标轮廓%d个" % len(cnt))
        cnt = pre_object_detect(cnt, width, height, ra=[0.01, 0.8], rw=[0.08, 0.8], rh=[0.03, 0.8])
        cnt = counter_similarity(cnt)
        cv.drawContours(contour_origin, cnt, -1, (0, 0, 255), 2)
        cv.imwrite(save_path + "/str_contour.jpg", contour_origin)
        print("剩下直行车道地标轮廓%d个" % len(cnt))
        return cnt[0]
    elif templte == 1 or templte == "left":
        # 读取模板
        img = cv.imread("./templte/left.png")
        # img = cv.resize(img, (width, height))
        # 边缘检测
        edge = cv.Canny(img, 50, 150, apertureSize=3)
        cv.imwrite(save_path + "/left_edge.jpg", edge)
        # 二值图像提取轮廓，轮廓跟踪算法（Suzuki，1985）
        contours, hierarchy = cv.findContours(edge, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
        contour_origin_pre = img.copy()
        contour_origin = img.copy()
        cnt = contours[0:len(contours)]
        cv.drawContours(contour_origin_pre, cnt, -1, (0, 0, 255), 2)
        cv.imwrite(save_path + "/left_contour_pre.jpg", contour_origin_pre)
        print("预检左转车道地标轮廓%d个" % len(cnt))
        cnt = pre_object_detect(cnt, width, height, ra=[0.01, 0.8], rw=[0.08, 0.8], rh=[0.03, 0.8])
        cnt = counter_similarity(cnt)
        cv.drawContours(contour_origin, cnt, -1, (0, 0, 255), 2)
        cv.imwrite(save_path + "/left_contour.jpg", contour_origin)
        print("剩下左转车道地标轮廓%d个" % len(cnt))
        return cnt[0]


def lane_direction_detect(img, mask):
    # 尺寸变化
    img = cv.resize(img, (width, height))
    # 原图
    origin = img
    # ROI提取
    ROI = origin.copy()
    ROI = cv.bitwise_and(mask, ROI)
    ROI = cv.cvtColor(ROI, cv.COLOR_BGR2GRAY)
    # ROI[0:int(height * 0.4), :] = 0
    # ROI[int(height * 0.8):height, :] = 0
    cv.imwrite(save_path + "/roi.jpg", ROI)
    # 边缘检测
    edge = cv.Canny(ROI, 50, 150, apertureSize=3)
    cv.imwrite(save_path + "/edge_roi.jpg", edge)
    # 二值化
    # binary = cv.adaptiveThreshold(edge, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 11, 5)
    # binary = cv.bitwise_not(binary)
    # cv.imwrite(save_path + "/binary_roi.jpg", binary)
    # 形态学操作
    # element = cv.getStructuringElement(cv.MORPH_RECT, (1, 1))
    # element2 = cv.getStructuringElement(cv.MORPH_RECT, (1, 1))
    # 开运算
    # morphology = cv.morphologyEx(edge, cv.MORPH_OPEN, element)
    # 膨胀运算
    # morphology = cv.dilate(edge, element2)
    # cv.imwrite(save_path + "/morphology.jpg", morphology)
    # 二值图像提取轮廓，轮廓跟踪算法（Suzuki，1985）
    contours, hierarchy = cv.findContours(edge, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
    contour_origin_pre = origin.copy()
    contour_origin = origin.copy()
    cnt = contours[0:len(contours) - 2]
    cv.drawContours(contour_origin_pre, cnt, -1, (0, 0, 255), 2)
    cv.imwrite(save_path + "/contour_pre.jpg", contour_origin_pre)
    print("预检车道地标轮廓%d个" % len(cnt))
    cnt = pre_object_detect(cnt, width, height)
    cnt = counter_similarity(cnt)
    print("剩下车道地标轮廓%d个" % len(cnt))
    # cv.drawContours(contour_origin, cnt, -1, color=(0, 255, 0), thickness=cv.FILLED)
    cv.drawContours(contour_origin, cnt, -1, color=(0, 255, 0), thickness=2)
    cv.imwrite(save_path + "/contour.jpg", contour_origin)
    # 模板匹配
    direction = contour_origin.copy()
    stra = templte_read("straight")
    left = templte_read("left")
    print(len(cnt))
    for i in range(len(cnt)):
        M = cv.moments(cnt[i])
        cX = int(M["m10"] / M["m00"])
        cY = int(M["m01"] / M["m00"])
        print(len(stra), len(cnt[i]))
        ret1 = cv.matchShapes(stra, cnt[i], 3, 0.0)
        ret2 = cv.matchShapes(left, cnt[i], 3, 0.0)
        if ret1 > ret2:
            print("轮廓%d: 直行" % (i))
            cv.putText(direction, "Straight", (cX - 40, cY - 40),
                        cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        else:
            print("轮廓%d: 左转" % (i))
            cv.putText(direction, "Left", (cX - 40, cY - 40),
                       cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        # print("轮廓%d: 直行匹配度: %.3f,左转匹配度: %.3f" % (i, ret1, ret2))
    cv.imwrite(save_path + "/contour_direction.jpg", direction)

    # Harris角点检测
    harris = origin.copy()
    mask = cv.bitwise_and(mask, harris)
    gray = cv.cvtColor(mask, cv.COLOR_BGR2GRAY)
    ROI = gray.copy()
    ROI[0:int(height * 0.30), :] = 0
    cv.imwrite(save_path + "/harris-roi.jpg", ROI)
    # cornerHarris函数图像格式为 float32
    ROI = np.float32(ROI)
    J = (0.05, 0.01, 0.005, 0.0005)
    for j in J:  # 遍历设置阈值：j * dst.max()
        dst = cv.cornerHarris(src=ROI, blockSize=5, ksize=7, k=0.04)
        a = dst > j * dst.max()
        harris[a] = [0, 255, 0]
        cv.imwrite(save_path + "/" + "corners_" + str(j) + ".jpg", harris)
        # cv.imshow('corners_' + str(j), img)

    # Shi-Tomasi角点检测
    tomasi = origin.copy()
    mask = cv.bitwise_and(mask, tomasi)
    gray = cv.cvtColor(mask, cv.COLOR_BGR2GRAY)
    ROI = gray.copy()
    ROI[0:int(height * 0.30), :] = 0
    # ROI[int(height * 0.82):height, :] = 0
    # ROI[:, int(width * 0.85):width] = 0
    keypoints = getkpoints(ROI, ROI)
    if keypoints is not None and len(keypoints) > 0:
        for x, y in keypoints:
            cv.circle(tomasi, (int(x), int(y)), 3, (0, 0, 255))
    cv.imwrite(save_path + "/Shi-Tomasi.jpg", tomasi)

    # 亚像素检测
    subpixel = origin.copy()
    mask = cv.bitwise_and(mask, tomasi)
    gray = cv.cvtColor(mask, cv.COLOR_BGR2GRAY)
    ROI = gray.copy()
    ROI[0:int(height * 0.30), :] = 0

    dst = cv.cornerHarris(ROI, 5, 7, 0.04)
    # dst = cv.dilate(dst, None)
    # ret, dst = cv.threshold(dst, 0.01 * dst.max(), 255, 0)
    dst = np.uint8(dst)
    ret, labels, stats, centroids = cv.connectedComponentsWithStats(dst)
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 300, 0.1)
    corners = cv.cornerSubPix(ROI, np.float32(centroids), (3, 3), (-1, -1), criteria)
    # 将检测到的亚像素角点绘制到原图上
    for i in range(corners.shape[0]):
        cv.circle(subpixel, (int(centroids[i][0]), int(centroids[i][1])), 3, (255, 0, 0))

    # # Now draw them
    # res = np.hstack((centroids, corners))
    # # np.int0 可以用来省略小数点后的数字，非四舍五入
    # res = np.int0(res)
    # subpixel[res[:, 1], res[:, 0]] = [0, 0, 255]
    # subpixel[res[:, 3], res[:, 2]] = [0, 255, 0]
    cv.imwrite(save_path + "/subpixel.jpg", subpixel)


if __name__ == '__main__':
    img_path = "./picture/263.jpg"
    if os.path.exists(img_path):
        img = cv.imread(img_path)
        bw_mask, color_mask = lane_line_detect(img)
        lane_direction_detect(img, bw_mask)
    else:
        print("图片不存在")
